#!/bin/bash

source collect/collectiontool.sh
source stress/stress_test.sh

declare -g collection_dir="/tmp/collect"

function generate_result_dir() {
    if [[ -e ${collection_dir} ]]; then
        mv ${collection_dir} ${collection_dir}_$(date +%Y_%m_%d_%H_%M_%S)
    fi

    mkdir -p ${collection_dir}
}

function usage() {
    echo "Usage: $0 [options]"
    echo "Run collection and stress test"
    echo "Options:"
    echo "    -c, --collect                     Enable machine and container info collection"
    echo "    -t, --timeout                     Collection timeout time"
    echo ""
    echo "    -s, --stress                      Enable stress test"
    echo "        --stress-collect              Enable stress info collection"
    echo "        --stress-cpu                  Enable cpu stress"
    echo "        --stress-memory               Enable memory stress"
    echo "        --stress-memory-bandwidth     Enable memory bandwidth stress"
    echo "        --stress-cache                Enable cache stress"
    echo "        --stress-cpu-cache-mixed      Enable cpu and cache mixed stress"
    echo ""
    echo "    -w, --stress-wait                 Run stress after wait time"
    echo "        --disk-location               Enable and run disk io stress at location"
    echo "        --network-ip                  Enable and run network stress connect to server IP"
    echo "        --network-port                Network stress connect to server port(default 5201)"
    echo ""
    echo "    -h, --help                        Help information"
}

function cmd_timeout {
    local cmd="$1"
    local timeout="$2";

    eval "$cmd" &
    local child=$!
    trap -- "" SIGTERM
    (
        sleep $timeout
        kill $child 2> /dev/null
    ) &
    wait $child
}

args=$(getopt -o cf:sf:t:w:h --long collect,stress,stress-collect,stress-cpu,stress-memory,stress-memory-bandwidth,stress-cache,stress-wait:,disk-location:,network-ip:,network-port:,stress-cpu-cache-mixed,help -- "$@")
if [ $? != 0 ] ; then echo "Terminating..." >&2 ; exit 1 ; fi
eval set -- "$args"

while true; do
    case "${1}" in
        -c|--collect)               collect_enable="true" ; shift ;;
        -t|--timeout)               timeout=$2 ; shift 2 ;;

        -s|--stress)                stress_enable="true" ; shift ;;
        --stress-collect)           stress_collect="true" ; shift ;;
        --stress-cpu)               stress_cpu="true" ; shift ;;
        --stress-memory)            stress_memory="true" ; shift ;;
        --stress-memory-bandwidth)  stress_memory_bandwidth="true" ; shift ;;
        --stress-cache)             stress_cache="true" ; shift ;;
        --stress-cpu-cache-mixed)   stress_cpu_cache_mixed="true" ; shift ;;

        -w|--stress-wait)           stress_wait=$2 ; shift 2 ;;
        --disk-location)            disk_location=$2 ; shift 2 ;;
        --network-ip)               network_ip=$2 ; shift 2 ;;
        --network-port)             network_port=$2 ; shift 2 ;;

        -h|--help)                  usage ; exit 0 ;;
        --)                         shift ; break ;;
        *)                          echo "invalid parameter" ; exit -1 ;;
    esac
done

generate_result_dir

if [ "x$collect_enable" == "xtrue" ]; then
    if [ "x$timeout" != "x" ]; then
        cmd_timeout "collect_machine_loop" $timeout &
        cmd_timeout "collect_container_loop" $timeout &

    else
        collect_machine_loop &
        collect_container_loop &
    fi
fi

if [ "x$stress_enable" == "xtrue" ]; then
    stress_checker
    if [ $? -ne 0 ];  then
        echo "stress test failed without tool \"$stress_tool\""
        exit 1
    fi

    if [ "x$stress_collect" == "xtrue" ]; then
        if [ "x$timeout" != "x" ]; then
            cmd_timeout "collect_stress_usage_loop" $timeout &
        else
            collect_stress_usage_loop &
        fi
    fi

    if [ "x$stress_wait" != "x" ]; then
        sleep $stress_wait
    fi

    if [ "x$stress_cpu" == "xtrue" ]; then
        stress_cpu
    fi
    
    if [ "x$stress_memory" == "xtrue" ]; then
        stress_memory
    fi

    if [ "x$stress_cache" == "xtrue" ]; then
        stress_cache
    fi

    if [ "x$disk_location" != "x" ]; then
        stress_disk_io $disk_location
    fi

    if [ "x$stress_memory_bandwidth" == "xtrue" ]; then
        stress_memory_bandwidth
    fi

    if [ "x$stress_cpu_cache_mixed" == "xtrue" ]; then
        stress_cpu_cache_mixed
    fi
fi

if [ "x$network_ip" != "x" ]; then
    if [ "x$network_port" == "x" ]; then
        network_port=5201
    fi

    # network stress test need two node
    stress_network $network_ip $network_port
fi
